#include "kernel_operator.h"
// #include "kernel_float.h"
// #include "impl/"
#include <cstdint>
using namespace AscendC;
// #define K_MAX_SHAPE_DIM 0

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue
template<typename T>class KernelArgMaxValuePerf{
public:
    __aicore__ inline KernelArgMaxValuePerf(){}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indice, GM_ADDR values,
     // uint32_t totalLength, uint32_t resLength, uint32_t stride,
     // uint32_t iterStep, uint32_t nIter, int64_t axes, uint32_t dims, 
     uint32_t axesDim,
     uint32_t smallDataNum,uint32_t bigDataNum,uint32_t tailBlockNum, TPipe* pipeIn){

        this->axesDim = axesDim;
        // 只处理round到512的
        this->roundDim = 512;
        pipe = pipeIn; 
        uint32_t coreNum = AscendC::GetBlockIdx();
        uint32_t globalIndex = bigDataNum * coreNum;
        if (coreNum < tailBlockNum) {
            this->coreDataNum = bigDataNum;
        } else {
            this->coreDataNum = smallDataNum;
            globalIndex -= (bigDataNum - smallDataNum)*(coreNum - tailBlockNum);
        }
        this->roundcoreDataNum = (this->coreDataNum + 7)/8*8;
        // printf("global: %d coreDataNum:%d\n", globalIndex, coreDataNum);
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + globalIndex*axesDim, this->coreDataNum*axesDim *sizeof(DTYPE_X));
        indiceGm.SetGlobalBuffer((__gm__ DTYPE_INDICE*)indice + globalIndex, this->coreDataNum * sizeof(DTYPE_INDICE));
        valuesGm.SetGlobalBuffer((__gm__ DTYPE_VALUES*)values + globalIndex, this->coreDataNum * sizeof(DTYPE_VALUES));
        
        // this->loopCount = this->coreDataNum;
        pipe->InitBuffer(full, this->roundDim*this->coreDataNum*sizeof(DTYPE_X) + 5*this->roundcoreDataNum*sizeof(DTYPE_X));
                         

    }
    __aicore__ inline void process() {
        copyIn();
        compute();
        // copyOut();
    }
    __aicore__ inline void copyIn() {
        xLocal = full.GetWithOffset<DTYPE_X>(this->roundDim*this->coreDataNum, 0);
        AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(coreDataNum),  static_cast<uint32_t>(axesDim*4), 0,  (64-(axesDim+7)/8),0};
        AscendC::DataCopyPadExtParams<DTYPE_X> padParams{true, 0,  static_cast<uint8_t>(8 - axesDim%8), DTYPE_X(-1000)};
        AscendC::DataCopyPad(xLocal, xGm, copyParams,padParams);
        xLocal.GetValue(0);
    }
    __aicore__ inline void compute() {
        // auto xLocal = full.GetWithOffset<DTYPE_X>(this->roundDim*this->coreDataNum, 0);
        AscendC::LocalTensor<int32_t> index = full.GetWithOffset<int32_t>(this->roundcoreDataNum, this->roundDim*this->coreDataNum*4);
        AscendC::LocalTensor<DTYPE_X> workLocal =full.GetWithOffset<DTYPE_X>(this->roundcoreDataNum, (this->roundDim*this->coreDataNum + this->roundcoreDataNum)*4);
        AscendC::LocalTensor<DTYPE_X> workLocal2 = full.GetWithOffset<DTYPE_X>(this->roundcoreDataNum, (this->roundDim*this->coreDataNum + this->roundcoreDataNum*2)*4);
        AscendC::LocalTensor<DTYPE_X> valueLocal = full.GetWithOffset<DTYPE_X>(this->roundcoreDataNum, (this->roundDim*this->coreDataNum + this->roundcoreDataNum*3)*4);
         AscendC::LocalTensor<int32_t> indexLocal = full.GetWithOffset<int32_t>(this->roundcoreDataNum, (this->roundDim*this->coreDataNum + this->roundcoreDataNum*4)*4);
        AscendC::WholeReduceMax<DTYPE_X>(xLocal, xLocal, 64, coreDataNum*8, 1, 1, 8, AscendC::ReduceOrder::ORDER_VALUE_INDEX);
        auto workLocalInt = workLocal.ReinterpretCast<int32_t>();
        AscendC::CreateVecIndex(workLocalInt, (int32_t)0, coreDataNum);
        AscendC::Muls(workLocalInt, workLocalInt, int32_t(16) , coreDataNum);
        
        AscendC::WholeReduceMax<DTYPE_X>(workLocal2, xLocal, 16, coreDataNum, 1, 1, 2, AscendC::ReduceOrder::ORDER_ONLY_INDEX);

        AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(coreDataNum * 4), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
//         valueLocal.GetValue(0);

        auto workLocal2Int = workLocal2.ReinterpretCast<int32_t>();
        // auto workLocalInt = workLocal.ReinterpretCast<int32_t>();
        AscendC::Add(workLocalInt,workLocalInt, workLocal2Int, coreDataNum);
        AscendC::Muls(workLocalInt, workLocalInt, int32_t(4) , coreDataNum);
        Gather(valueLocal, xLocal, workLocalInt.ReinterpretCast<uint32_t>(), (uint32_t)0,coreDataNum);
        AscendC::Adds(workLocalInt, workLocalInt, int32_t(4) , coreDataNum);
        Gather(indexLocal, xLocal.ReinterpretCast<int32_t>(), workLocalInt.ReinterpretCast<uint32_t>(), (uint32_t)0,coreDataNum);

        AscendC::Muls(workLocal2Int, workLocal2Int, int32_t(32) , coreDataNum);
        AscendC::Add(indexLocal,indexLocal, workLocal2Int, coreDataNum);
        indexLocal.GetValue(0);
        // for (int i = 0; i < 5000; i++) {
        // indexLocal.GetValue(0);
        // }
        AscendC::DataCopyPad(indiceGm, indexLocal, copyParams); // 从GM->VECIN搬运40Bytes
        AscendC::DataCopyPad(valuesGm, valueLocal, copyParams); // 从GM->VECIN搬运40Bytes

    }


private:
    int32_t coreDataNum;
    int64_t axes;
    uint32_t axesDim;
    uint32_t roundDim,roundcoreDataNum;
    // uint32_t loopCount;
    AscendC::TPipe *pipe;
    AscendC::GlobalTensor<DTYPE_X> xGm;
    AscendC::GlobalTensor<DTYPE_INDICE> indiceGm;
    AscendC::GlobalTensor<DTYPE_VALUES> valuesGm;
    TBuf<QuePosition::VECCALC>full;
    AscendC::LocalTensor<DTYPE_X>xLocal;
};
